import pandas as pd
import streamlit as st

from agent.csv_agent_data import get_data_from_csv


def create_chart(input_data, chart_type):
    df_data = pd.DataFrame(input_data["data"], columns=input_data["columns"])
    df_data.set_index(input_data["columns"][0], inplace=True)
    if chart_type == "table":
        st.table(df_data)
    elif chart_type == "bar":
        st.bar_chart(df_data)
    elif chart_type == "line":
        st.line_chart(df_data)
    elif chart_type == "scatter":
        st.line_chart(df_data)


with st.sidebar:
    openai_key = st.text_input("输入OpenAI API Key", type="password")

st.header("CSV智能数据分析")
df = st.file_uploader("上传CSV文件", type=["csv"])
if df:
    st.session_state["df"] = pd.read_csv(df)
    with st.expander("展示数据"):
        st.dataframe(st.session_state["df"])

question = st.text_input("请根据上面提交的csv文件，进行问题提问(支持简单文本，表格，折线图，条形图，散点图)的显示")
flag = st.button("提交问题")

if flag:
    if 'df' in st.session_state:
        if st.session_state['df'].empty:
            st.warning("请先上传CSV文件")
            st.stop()
    else:
        st.warning("请先上传CSV文件")
        st.stop()
    if not question:
        st.warning("请输⼊问题")
        st.stop()
    if not openai_key:
        st.info("请输⼊你的Openai密钥")
        st.stop()
    with st.spinner("正在分析数据"):
        response = get_data_from_csv(openai_key, st.session_state['df'], question)
        if "answer" in response:
            st.write(response["answer"])
        if "table" in response:
            create_chart(response["table"], "table")
        if "bar" in response:
            create_chart(response["bar"], "bar")
        if "line" in response:
            create_chart(response["line"], "line")
        if "scatter" in response:
            create_chart(response["scatter"], "scatter")
